"""
Contains some simple functions for working with the Safarzadeh et al. (2015)
FIR SED templates.

The find_template function returns the SED template that
is the best match for the log (L_IR/L_sun) and log (M_dust/M_sun) values
specified by the user.

The load_templates function reads the ASCII file that
contains the template data and returns arrays that contain the L_IR and M_dust
values for each template, a 2-D array that contains the normalized SEDs
(lambda*L_lambda/L_IR), and the wavelength array.

"""

import numpy as np

def find_template(
    lir, #"""logarithm of the IR luminosity in solar units"""
    mdust, #"""logarithm of the dust mass in solar units"""
    tol=0.5, #"""maximum allowed discrepancy between template and requested values, delta(L_IR, M_dust); see function docstring"""
    rtol=0.2, #"""maximum allowed difference between the template and requested L_IR/M_dust ratio"""
    sed_file="./Safarzadeh_et_al_2015_FIR_SED_templates.txt", #"""full path to file that contains the templates"""
    verbose=True #"""prints messages if set to True"""
    ):
    """Return template with (L_IR, M_dust) values closest to those requested.

    Similarity is defined by the Euclidean distance between the requested and template 
    (L_IR, M_dust) values,

        delta(L_IR, M_dust) = [(L_IR,requested - L_IR, template)^2
            + (M_dust,requested - M_dust,template)^2]^(0.5)

    If no template with delta(L_IR, M_dust) < tol is found, then search for a template with an
    L_IR/M_dust ratio that differs from that requested by at most rtol.

    Parameters:
        - lir : log (L_IR/L_sun) value for which a template is desired
        - mdust : log (M_dust/M_sun) value for which a template is desired
        - tol : maximum allowed delta(L_IR, M_dust) value
        - rtol : maximum allowed difference between requested and template L_IR/M_dust ratio;
        only used if delta(L_IR, M_dust) < tol cannot be satisfied
        - sed_file : full path to file that contains the template
        - verbose : print helpful messages if set to True

    Returns:
        - lambda_array: array of wavelengths (in micron)
        - sed: template SED that is most appropriate for the (L_IR, M_dust) requested; the
            template is normalized such that its total IR luminosity equals  the *requested*
            L_IR value (i.e. the normalized SED from the template file is multipled by L_IR

    Example:
        Find an SED for a galaxy with L_IR = 10^11 L_sun and M_dust = 10^8 M_sun:

            lam, sed = find_template(11.,8.)
"""


    lir_array, mdust_array, sed_array, lambda_array = load_templates()

    dist = ((lir_array - lir)**2 + (mdust_array - mdust)**2)**(0.5)
    
    # first try to find a template with delta(L_IR,M_dust) < tol (see function docstring)
    if dist.min() < tol :
        if verbose :
            print "Template sufficiently close in the (L_IR, M_dust) plane found"
            print "Requested: log L_IR = ", lir, ", log M_dust = ", mdust
            print "Template: log L_IR = ", lir_array[dist.argmin()], ", log M_dust = ", \
                mdust_array[dist.argmin()]
        return lambda_array, sed_array[dist.argmin(),:]*lir
    else : # not found, so search for one with delta(L_IR/M_dust) < rtol (see function docstring)
        if verbose :
            print "Template with delta(L_IR, M_dust) < ",tol," not found. Returning template with nearest L_IR/M_dust value instead"
        ratio=lir-mdust
        ratio_array=lir_array-mdust_array
        dist_ratio = np.abs(ratio - ratio_array)
        if dist_ratio.min() < rtol :
            if verbose :
                print "Found template with sufficiently similar L_IR/M_dust ratio"
                print "Requested: log L_IR = ", lir, ", log M_dust = ", mdust, \
                    ", log L_IR/M_dust = ",ratio
                print "Template: log L_IR = ", lir_array[dist_ratio.argmin()], ", log M_dust = ", \
                    mdust_array[dist_ratio.argmin()],", log L_IR/M_dust = ", \
                    ratio_array[dist_ratio.argmin()]
            sed = sed_array[dist_ratio.argmin(),:]*lir
            return lambda_array, sed
        else : # didn't find a suitable template
            print "ERROR: acceptable template not found"
            return
        
def load_templates(
    sed_file = "./Safarzadeh_et_al_2015_FIR_SED_templates.txt" #"""full path to file that contains the templates"""
    ):
    """Read the ASCII file that contains the template data.

    Parameters:
        - sed_file: the full path to the template file.

    Returns:
        - lir_array : log (L_IR/L_sun) for each template

        - mdust_array : log (M_dust/M_sun) for each template

        - sed_array: the actual SEDs (lambda*L_lambda) normalized by dividing by L_IR.
        The first dimension is the template number, and the second is wavelength.
    
        - lambda_array: the wavelengths (in microns) at which the SEDs are provided

    Example:
        import template_functions as tf
        lir_array, mdust_array, sed_array, lambda_array = tf.load_templates()
    """

    lir_array = np.loadtxt(sed_file,skiprows=21,usecols=[0])
    mdust_array = np.loadtxt(sed_file,skiprows=21,usecols=[1])
    sed_array = np.loadtxt(sed_file,skiprows=21,usecols=(np.arange(19)+3))
    lambda_array = np.loadtxt(sed_file,skiprows=20,usecols=(np.arange(19)+3))[0,:]*1.e6

    return lir_array, mdust_array, sed_array, lambda_array
